#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include <cufft.h>
#include <time.h>
#include "prof.h"

/* CUDA错误处理宏 */
#define CHECK_CUDA_ERROR(call) { \
    cudaError_t err = call; \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA错误: %s\n", cudaGetErrorString(err)); \
        exit(EXIT_FAILURE); \
    } \
}

/* cuFFT错误处理宏 */
#define CHECK_CUFFT_ERROR(call) { \
    cufftResult err = call; \
    if (err != CUFFT_SUCCESS) { \
        fprintf(stderr, "cuFFT错误: %d\n", err); \
        exit(EXIT_FAILURE); \
    } \
}

int main(int argc, char** argv) {
    int n, m;

    // 获取运行次数n
    if (argc != 2) {
        fprintf(stderr, "用法: %s <运行次数>\n", argv[0]);
        return 1;
    }
    n = atoi(argv[1]);
    if (n <= 0) {
        fprintf(stderr, "运行次数必须为正整数。\n");
        return 1;
    }

    profInit();//初始化
    
    profSetSectName(1, "stft 实数 htod 时间");
    profSetSectName(2, "stft 实数 compute_time");
    profSetSectName(3, "stft 实数 dtoh 时间");

    profSetSectName(4, "stft 复数 htod 时间");
    profSetSectName(5, "stft 复数 compute_time");
    profSetSectName(6, "stft 复数 dtoh 时间");

    /* 定义参数 */
    int W = 8192;      // 窗口大小
    int K = 10;        // 每个信号的窗口数
    int N = W * K;     // 信号长度
    // 组数（m）：动态输入，表示并行计算的信号组数。
    m = 1000;

    /* 为实数STFT分配内存 */
    float *h_signal_real;    // 主机端实数信号
    float *d_signal_real;    // 设备端实数信号
    cufftComplex *d_fft_real; // 设备端FFT结果
    cufftComplex *h_fft_real; // 主机端FFT结果
    h_signal_real = (float*)malloc(m * N * sizeof(float));
    CHECK_CUDA_ERROR(cudaMalloc(&d_signal_real, m * N * sizeof(float)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_fft_real, m * K * (W / 2 + 1) * sizeof(cufftComplex)));
    h_fft_real = (cufftComplex*)malloc(m * K * (W / 2 + 1) * sizeof(cufftComplex));

    /* 为复数STFT分配内存 */
    cufftComplex *h_signal_complex; // 主机端复数信号
    cufftComplex *d_signal_complex; // 设备端复数信号
    cufftComplex *d_fft_complex;    // 设备端FFT结果
    cufftComplex *h_fft_complex;    // 主机端FFT结果
    h_signal_complex = (cufftComplex*)malloc(m * N * sizeof(cufftComplex));
    CHECK_CUDA_ERROR(cudaMalloc(&d_signal_complex, m * N * sizeof(cufftComplex)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_fft_complex, m * K * W * sizeof(cufftComplex)));
    h_fft_complex = (cufftComplex*)malloc(m * K * W * sizeof(cufftComplex));

    /* 生成随机输入数据 */
    srand(time(NULL));
    for (int i = 0; i < m * N; i++) {
        h_signal_real[i] = (float)rand() / RAND_MAX; // 实数信号，范围[0,1]
        h_signal_complex[i].x = (float)rand() / RAND_MAX; // 复数实部
        h_signal_complex[i].y = (float)rand() / RAND_MAX; // 复数虚部
    }

    /* 创建cuFFT计划 */
    cufftHandle plan_real;    // 实数STFT计划
    CHECK_CUFFT_ERROR(cufftPlanMany(&plan_real, 1, &W, NULL, 1, W, NULL, 1, W / 2 + 1, CUFFT_R2C, m * K));

    cufftHandle plan_complex; // 复数STFT计划
    CHECK_CUFFT_ERROR(cufftPlanMany(&plan_complex, 1, &W, NULL, 1, W, NULL, 1, W, CUFFT_C2C, m * K));

    profEnterSect(0);
    /* 测试实数STFT */
    printf("\n运行实数STFT...\n");
    for (int i = 0; i < n; i++) {

        /* 记录htod时间 */
        profEnterSect(1);
        CHECK_CUDA_ERROR(cudaMemcpy(d_signal_real, h_signal_real, m * N * sizeof(float), cudaMemcpyHostToDevice));
        profLeaveSect(1);

        /* 记录FFT时间 */
        profEnterSect(2);
        CHECK_CUFFT_ERROR(cufftExecR2C(plan_real, d_signal_real, d_fft_real));
        CHECK_CUDA_ERROR(cudaDeviceSynchronize());
        profLeaveSect(2);

        /* 记录dtoh时间 */
        profEnterSect(3);
        CHECK_CUDA_ERROR(cudaMemcpy(h_fft_real, d_fft_real, m * K * (W / 2 + 1) * sizeof(cufftComplex), cudaMemcpyDeviceToHost));
        profLeaveSect(3);

    }

    /* 测试复数STFT */
    printf("\n运行复数STFT...\n");
    for (int i = 0; i < n; i++) {

        /* 记录htod时间 */
        profEnterSect(4);
        CHECK_CUDA_ERROR(cudaMemcpy(d_signal_complex, h_signal_complex, m * N * sizeof(cufftComplex), cudaMemcpyHostToDevice));
        profLeaveSect(4);

        /* 记录FFT时间 */
        profEnterSect(5);
        CHECK_CUFFT_ERROR(cufftExecC2C(plan_complex, d_signal_complex, d_fft_complex, CUFFT_FORWARD));
        CHECK_CUDA_ERROR(cudaDeviceSynchronize());
        profLeaveSect(5);

        /* 记录dtoh时间 */
        profEnterSect(6);
        CHECK_CUDA_ERROR(cudaMemcpy(h_fft_complex, d_fft_complex, m * K * W * sizeof(cufftComplex), cudaMemcpyDeviceToHost));
        profLeaveSect(6);

    }

    profLeaveSect(0);
    profPrintInfo();

    /* 输出统计报告 */
    printf("\n=== 统计报告 ===\n");

    /* 清理资源 */
    free(h_signal_real);
    free(h_signal_complex);
    free(h_fft_real);
    free(h_fft_complex);
    CHECK_CUDA_ERROR(cudaFree(d_signal_real));
    CHECK_CUDA_ERROR(cudaFree(d_signal_complex));
    CHECK_CUDA_ERROR(cudaFree(d_fft_real));
    CHECK_CUDA_ERROR(cudaFree(d_fft_complex));
    CHECK_CUFFT_ERROR(cufftDestroy(plan_real));
    CHECK_CUFFT_ERROR(cufftDestroy(plan_complex));

    return 0;
}
